# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np
import os
import shutil
from multiprocessing import Process
import time

"""
    并行计算气动力初始化
"""

# 开启进程数
split_nums = 5

# 数据集名称
dataset = "CST_20"

# 待计算气动力的数据集
dataset_dir = r"/home/linjx/data/dataset/UIUC_AND_NACA_" + dataset


def Init():
    """
        初始化目录
        在当前的目录下新建一个para目录
        并在para目录下新建{$split_nums}个算例，命名为stdsample{$i}
    :return:
    """
    if os.path.exists("./para"):
        shutil.rmtree("./para")
    os.mkdir("./para")

    for i in range(1, split_nums + 1):
        if os.path.exists("./para/stdsample{}".format(i)):
            shutil.rmtree("./para/stdsample{}".format(i))
        shutil.copytree("./stdsample_para", "./para/stdsample{}".format(i))


def split(dataset_dir=r"/home/linjx/data/dataset/UIUC_AND_NACA_CST_20"):
    """
        读出所有的翼型文件名称，并将其进行均匀划分
    :param dataset_dir:
    :return:
    """
    names_list = []
    for file_name in os.listdir(dataset_dir):
        names_list.append(file_name)

    pad = len(names_list) // split_nums

    # split
    for i in range(split_nums - 1):
        divide_i = np.asarray(names_list[i * pad:(i + 1) * pad]).reshape((1, -1))
        df = pd.DataFrame(divide_i)
        df.to_csv("./para/stdsample{}/solver/names.csv".format(i + 1),
                  index=False,
                  header=False)

    # 最后一个进程单独划分
    divide_i = np.asarray(names_list[(split_nums - 1) * pad::]).reshape((1, -1))
    df = pd.DataFrame(divide_i)
    df.to_csv("./para/stdsample{}/solver/names.csv".format(split_nums),
              index=False,
              header=False)


def worker(i):
    """
        执行第i个子进程
    :param i:
    :return:
    """
    print("子进程执行中>>> pid={0},ppid={1}".format(os.getpid(), os.getppid()))

    # 切换到工作目录
    work_dir = os.path.join(os.getcwd(), "para/stdsample{}/solver/".format(i))
    os.chdir(work_dir)

    # 执行foilFlow程序
    foilFlow_path = os.path.join(os.getcwd(),
                                 "para/stdsample{}/solver/loop_foilFlow.py".format(i))
    os.system(
        "nohup python {} -d {} &".format(foilFlow_path, dataset))

    print("子进程执行结束<<< pid={0},ppid={1}".format(os.getpid(), os.getppid()))


def run():
    """
        开启多进程
    :return:
    """
    print("主进程执行中>>> pid={0}".format(os.getpid()))
    ps = []

    # 创建子进程实例
    for i in range(1, split_nums + 1):
        p = Process(target=worker, name="worker" + str(i), args=(i,))
        ps.append(p)
        time.sleep(2)

    # 开启进程
    for i in range(split_nums):
        ps[i].start()

    # 阻塞进程（等待所有进程执行完毕）
    for i in range(split_nums):
        ps[i].join()

    print("主进程终止")


if __name__ == "__main__":
    Init()
    split(dataset_dir)
    run()
